##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(ggplot2)
library(ggsci)
library(ggpubr)

##########################################################################################

option_list <- list(
    make_option(c("--sample_list_file"), type = "character"),
    make_option(c("--igc_cluster_file"), type = "character"),
    make_option(c("--dgc_cluster_file"), type = "character"),
    make_option(c("--rsem_file"), type = "character"),
    make_option(c("--stem_file"), type = "character"),
    make_option(c("--stem_tpm_file"), type = "character"),
    make_option(c("--stem_plot_tpm"), type = "character"),
    make_option(c("--gtf_file"), type = "character"),
    make_option(c("--ddr_gene_file"), type = "character"),
    make_option(c("--mutsig_gene_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    sample_list_file <- "~/20220915_gastric_multiple/rna_combine/analysis/config/tumor_normal.list"
    igc_cluster_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/mfuzz_v2/mfuzz_plot_IGC.annotation.tsv"
    dgc_cluster_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/mfuzz_v2/mfuzz_plot_DGC.annotation.tsv"
    rsem_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/DiffGene/CombineCounts.FilterLowExpression-MergeMutiSample.TMM.tsv"
    stem_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/stemness/StemScore_Expression.tsv"
    stem_tpm_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/stemness/StemScore.MutipleStage.tsv"
    stem_plot_tpm <- "~/20220915_gastric_multiple/rna_combine/analysis/RSEM/CombineTpm.FilterLowExpression-MergeMutiSample.tsv"
    ddr_gene_file <- "~/20220915_gastric_multiple/rna_combine/analysis/public_ref/DDR.list"
    gtf_file <- "~/ref/GTF/gencode.v19.ensg_genename.txt"
    mutsig_gene_file <- "~/20220915_gastric_multiple/dna_combine/mutsig_check/smg.list"
    out_path <- "~/20220915_gastric_multiple/rna_combine/analysis/images/importGene_v2"

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_list_file <- opt$sample_list_file
igc_cluster_file <- opt$igc_cluster_file
dgc_cluster_file <- opt$dgc_cluster_file
out_path <- opt$out_path
rsem_file <- opt$rsem_file
stem_file <- opt$stem_file
stem_tpm_file <- opt$stem_tpm_file
ddr_gene_file <- opt$ddr_gene_file
gtf_file <- opt$gtf_file
stem_plot_tpm <- opt$stem_plot_tpm
mutsig_gene_file <- opt$mutsig_gene_file

dir.create(out_path , recursive = T)

##########################################################################################

info <- data.frame(fread(sample_list_file))
dat_igc_cluster <- data.frame(fread(igc_cluster_file))
dat_dgc_cluster <- data.frame(fread(dgc_cluster_file))
dat_expression <- data.frame(fread(rsem_file))
dat_stem <- data.frame(fread(stem_file))
dat_stem_tpm <- data.frame(fread(stem_tpm_file))
dat_stem_forplot <- data.frame(fread(stem_plot_tpm))[,-2]
ddr_gene_list <- data.frame(fread(ddr_gene_file , header = F))
dat_gtf <- data.frame(fread(gtf_file , header = F))
dat_smg <- data.frame(fread(mutsig_gene_file , header = T))

colnames(dat_stem) <- c("Hugo_Symbol" , "cor_stem" , "p_stem" , "q_stem" )
colnames(dat_gtf) <- c("gene_id" , "Hugo_Symbol")

##########################################################################################

igc_class <- c("IM + IGC + DGC" , "IM + IGC")
dgc_class <- c("IM + IGC + DGC" , "IM + DGC")

im_igc_sample <- unique(info[info$Type %in% igc_class , "ID"])
im_dgc_sample <- unique(info[info$Type %in% dgc_class , "ID"])

##########################################################################################
dat_expression <- merge( dat_expression , dat_gtf , by = "gene_id" )
dat_stem_forplot <- merge( dat_stem_forplot , dat_gtf , by = "gene_id" )

##########################################################################################
plotTpm <- function(tmp_dat = tmp_dat , out_name = out_name , title = title){

    my_comparisons_1 <- list( 
        c(1, 2), c(1, 3) , 
        c(2, 3) )

    plot <- ggplot( tmp_dat , aes( x = Class , y = Tpm , color = Class ) ) +
        geom_line( aes( group = Sample ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
        geom_boxplot(alpha =1 , outlier.size=0 , size = 0.9 , width = 0.6) +
        geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
        facet_grid(.~Type,space='free_x',scales='free_x') +
        scale_fill_npg()+
        scale_color_npg()+
        xlab(NULL) +
        ylab("TMM")+
        theme_bw() +
        ggtitle(title) +
        stat_compare_means(comparisons = my_comparisons_1) +
        theme(panel.background = element_blank(),#设置背影为白色#清除网格线
            legend.position ='none',
            legend.title = element_blank() ,
            panel.grid.major=element_line(colour=NA),
            plot.title = element_text(size = 8,color="black",face='bold'),
            legend.text = element_text(size = 10,color="black",face='bold'),
            axis.text.y = element_text(size = 10,color="black",face='bold'),
            axis.title.x = element_text(size = 10,color="black",face='bold'),
            axis.title.y = element_text(size = 10,color="black",face='bold'),
            axis.ticks.x = element_blank(),
            axis.text.x = element_text(size = 10,color="black",face='bold') ,
            axis.line = element_line(size = 0.5)) 

    ## 若表达极高则进行log10转化
    if(max(tmp_dat$Tpm) > 5000){
        plot <- plot + scale_y_log10()
    }
    ggsave(file=out_name,plot=plot,width=5,height=5)

}

plotStemTpm <- function(tmp_dat_stem = tmp_dat_stem , out_name = out_name){

    tmp_dat_stem <- unique(tmp_dat_stem[,c("ID" , "Sample" , "Class" , "Tpm" , "Ratio")])

    plot <- ggplot( tmp_dat_stem , aes( x = Tpm , y = Ratio ) ) +
        geom_point() + 
        geom_smooth(method = 'lm') +
        xlab("TPM")+
        ylab("Stem Ratio")+
        stat_cor(data=tmp_dat_stem, method = "spearman") +
        theme(panel.background = element_blank(),#设置背影为白色#清除网格线
            legend.position ='none',
            legend.title = element_blank() ,
            panel.grid.major=element_line(colour=NA),
            plot.title = element_text(size = 8,color="black",face='bold'),
            legend.text = element_text(size = 10,color="black",face='bold'),
            axis.text.y = element_text(size = 10,color="black",face='bold'),
            axis.title.x = element_text(size = 10,color="black",face='bold'),
            axis.title.y = element_text(size = 10,color="black",face='bold'),
            axis.ticks.x = element_blank(),
            axis.text.x = element_text(size = 10,color="black",face='bold') ,
            axis.line = element_line(size = 0.5)) 
    ggsave(file=out_name,plot=plot,width=4,height=4)

}

getPlotData <- function(tmp_exp = tmp_exp , im_igc_sample = im_igc_sample , im_dgc_sample = im_dgc_sample){

    sample <- sapply( strsplit(colnames(tmp_exp)[2:(ncol(tmp_exp)-1)] , "_") , "[" , 1)
    class <- sapply( strsplit(colnames(tmp_exp)[2:(ncol(tmp_exp)-1)] , "_") , "[" , 2)
    tpm <- as.numeric(tmp_exp[2:(ncol(tmp_exp)-1)])
    tmp_dat <- data.frame( Sample = sample , Class = class , Tpm = tpm )
    
    ## Normal->IM->IGC看基因的表达变化情况
    tmp_dat_igc <- subset( tmp_dat , Sample %in% im_igc_sample )
    tmp_dat_igc <- subset( tmp_dat_igc , Class %in% c( "Normal" , "IM" , "IGC" ))
    tmp_dat_igc$Class <- factor( tmp_dat_igc$Class , levels = c( "Normal" , "IM" , "IGC" ) , order = T )
    tmp_dat_igc$Type <- "IM + IGC"

    ## Normal->IM->DGC看基因的表达变化情况
    tmp_dat_dgc <- subset( tmp_dat , Sample %in% im_dgc_sample )
    tmp_dat_dgc <- subset( tmp_dat_dgc , Class %in% c( "Normal" , "IM" , "DGC" ))
    tmp_dat_dgc$Class <- factor( tmp_dat_dgc$Class , levels = c( "Normal" , "IM" , "DGC" ) , order = T )
    tmp_dat_dgc$Type <- "IM + DGC"

    ## Normal->IM->GC看基因的表达变化情况
    tmp_dat_gc <- tmp_dat
    tmp_dat_gc$Class <- ifelse( tmp_dat_gc$Class %in% c("IGC" , "DGC") , "GC" , tmp_dat_gc$Class )
    tmp_dat_gc$Class <- factor( tmp_dat_gc$Class , levels = c( "Normal" , "IM" , "GC" ) , order = T )
    tmp_dat_gc$Type <- "IM + GC"
    ## 一个既有IGC又有DGC，合并
    tmp_dat_gc <- tmp_dat_gc %>%
    group_by(Sample , Class) %>%
    summarize( Tpm = median(Tpm) , Type = unique(Type) )

    tmp_dat <- rbind( tmp_dat_igc , tmp_dat_dgc , tmp_dat_gc )
    tmp_dat$Type <- factor( tmp_dat$Type , levels = c("IM + GC" , "IM + IGC" , "IM + DGC") )

    return(tmp_dat)
}

describeGene <- function(level_use = level_use , dat_use = dat_use , dat_expression = dat_expression , image_path = image_path , level_name = level_name ){

    image_path <- paste0( out_path , "/" , level_name )
    dir.create(image_path , recursive = T)

    dat_use_tmp <- dat_use[level_use,]
    gene_list <- unique(dat_use_tmp$Hugo_Symbol)

    for(gene in gene_list){
        cluster <- unique(dat_use_tmp[dat_use_tmp$Hugo_Symbol==gene,"CLUSTER"])
        dat_use_tmp_2 <- dat_use_tmp[dat_use_tmp$Hugo_Symbol==gene,]
        hub_ratio <- paste0(round(unique(dat_use_tmp_2$Hub_membership_rank) , 3) * 100 , "%")

        title <- paste0(
            "Gene : " , gene , "\n" ,
            "Cluster : " , cluster , "\n" ,
            "Hub Gene :" , hub_ratio , "\n" , 
            "Stem Gene : " , length(which(dat_use_tmp_2$p_stem <0.05)) > 1 , "\n" ,
            "Report Gastric Driver : " , unique(dat_use_tmp_2$SMG) , "\t," , "CGC : " , unique(dat_use_tmp_2$Role.in.Cancer) , "\t," , "DDR : " , unique(dat_use_tmp_2$ddr) , "\n" , 
            gsub( "," , "\n" , unique(dat_use_tmp_2$pathway))
        )

        ## TMM
        tmp_exp <- subset( dat_expression , Hugo_Symbol == gene)
        tmp_dat <- getPlotData(tmp_exp = tmp_exp , im_igc_sample = im_igc_sample , im_dgc_sample = im_dgc_sample)
        out_name <- paste0( image_path , "/" , gene , ".pdf" )

        ## 画TMM在不同样本间的变化
        plotTpm(tmp_dat = tmp_dat , out_name = out_name , title = title)

        ## 该样本的表达和干细胞评分的关系
        ## TPM
        tmp_exp <- subset( dat_stem_forplot , Hugo_Symbol == gene)
        tmp_dat <- getPlotData(tmp_exp = tmp_exp , im_igc_sample = im_igc_sample , im_dgc_sample = im_dgc_sample)
        tmp_dat$ID <- paste0( tmp_dat$Sample , "_" , tmp_dat$Class )
        dat_stem_tpm$ID <- paste0( dat_stem_tpm$Sample , "_" , dat_stem_tpm$Class )

        tmp_dat_stem <- merge( tmp_dat , dat_stem_tpm[,c("ID" , "Ratio")] , by = "ID" )
        out_name <- paste0( image_path , "/" , gene , ".stem_tpm.pdf" )
        plotStemTpm(tmp_dat_stem = tmp_dat_stem , out_name = out_name)
    }

}

computeTPMChange <- function(dat_use = dat_use){

    ## Membership值也可以暗示两个向量之间的相关性。如果两个基因对于一个特定的cluster都有高的membership score，那么他们通常来说表达模式是相似的。
    ## 我们对于高于阈值α的基因，叫做这个cluster的α-core
    #dat_use <- subset( dat_use , MEM.SHIP > 0.6 )
    dat_use$ddr <- ifelse( dat_use$Hugo_Symbol %in% ddr_gene_list$V1 , "TRUE" , "FALSE" )

    ## 增加干细胞评分相关程度
    dat_use <- merge(dat_use , dat_stem , by = "Hugo_Symbol")

    ## 基因的重要程度排序
    ## SMG、DDR、CGC、HallMarks的通路数量、hub基因
    smg_gene <- which(dat_use$SMG=="TRUE")
    cgc_gene <- which(!(is.na(dat_use$Role.in.Cancer)) & dat_use$Role.in.Cancer!="" )
    ddr_gene <- which(dat_use$Hugo_Symbol %in% ddr_gene_list$V1)

    ## smg和CGC等价已报道
    report_gene <- unique(c(smg_gene , cgc_gene , ddr_gene))
    
    ## 通路基因
    pathway_gene <- which(dat_use$pathway_num >= 1)

    ## hub基因
    hub_gene <- which(dat_use$Hub_membership=="TRUE" )

    ## 表达与干细胞显著相关的基因
    stem_gene <- which(dat_use$q_stem <= 0.05)

    ## mutsig基因
    mutsig_gene <- which(dat_use$Hugo_Symbol %in% dat_smg$Gene_Symbol)

    ## 记录基因
    record_level <- table(c(report_gene , hub_gene , pathway_gene , stem_gene))
    # record_level <- table(c(report_gene , hub_gene , pathway_gene))

    ## 既为已报道、又在通路、又是hub、又与细胞干性相关
    level1 <- as.numeric(names(which(record_level == 4)))

    ## 为已报道、通路、hub的任三
    level2 <- as.numeric(names(which(record_level == 3)))

    ## 为已报道、通路、hub的任二
    level3 <- as.numeric(names(which(record_level == 2)))

    ## 为已报道、通路、hub的任一
    level4 <- as.numeric(names(which(record_level == 1)))

    ## 描绘基因在同一患者的改变情况
    level_use <- level1
    level_name <- paste0("level1_" , tumor)
    describeGene(level_use = level_use , dat_use = dat_use , dat_expression = dat_expression , image_path = image_path , level_name = level_name )

    level_use <- level2
    level_name <- paste0("level2_" , tumor)
    describeGene(level_use = level_use , dat_use = dat_use , dat_expression = dat_expression , image_path = image_path , level_name = level_name )

    ## level3和level4仅展示胃癌的显著突变基因
    level_use <- level3[level3 %in% mutsig_gene]
    level_name <- paste0("level3_" , tumor)
    describeGene(level_use = level_use , dat_use = dat_use , dat_expression = dat_expression , image_path = image_path , level_name = level_name )

    level_use <- level4[level4 %in% mutsig_gene]
    level_name <- paste0("level4_" , tumor)
    describeGene(level_use = level_use , dat_use = dat_use , dat_expression = dat_expression , image_path = image_path , level_name = level_name )

}



##########################################################################################
## 分单个基因，看在一个人多个样本的改变情况
dat_use <- dat_igc_cluster
tumor <- "IGC"
computeTPMChange(dat_use = dat_use)

dat_use <- dat_dgc_cluster
tumor <- "DGC"
computeTPMChange(dat_use = dat_use)

